#include "dnskv.h"

namespace dnskv {

static inline size_t ROUND(size_t x, size_t u) {
    return (x / u) * u;
}

thread_local static char local_buffer[kvalloc::SLAB_SIZE];
EpochBasedMemoryReclamationStrategy *ebr;

uint64_t miss = 0;

constexpr uint16_t SIMPLE_CRC = (uint16_t)0x5AB4;

void Serialize(ReqType op, _key_t k, const string & v, char * buffer) {
    Request * req = (Request *)buffer;
    req->rwlock.Clear(); // 将锁的初始位设定好
    req->op = op; 
    req->crc = SIMPLE_CRC; 
    req->val_size = v.size();
    *((uint64_t *)req->keyvalue) = k;
    memcpy(req->keyvalue + 8, v.c_str(), v.size());
}

bool Deserialize(char * buffer, std::string & out) {
    Request * req = (Request *)buffer;
    if(req->crc != SIMPLE_CRC) {
        return false;
    } else {
        out = std::string(req->keyvalue + 8, req->val_size);
        return true;
    }
}

DNSKV::DNSKV(const DNSOpt & opt) {
    if(!kvalloc::folder_exist(opt.dbdir.c_str())) {
        mkdir(opt.dbdir.c_str(), 0777);
    }
    if(!kvalloc::folder_exist(opt.nvmdir.c_str())) {
        mkdir(opt.nvmdir.c_str(), 0777);
    }
    hot_index_ = new morphtree::Morphtree;
    global_index_ = new tlbtree::TLBtree(opt.nvmdir + "/tlbtree.pool", opt.nvmsize);
    allocator_ = new kvalloc::KVAlloc(opt.dbdir + "/dnskv.db", opt.nvmdir + "/slabmeta.pool", opt.ssdsize / kvalloc::SLAB_SIZE);
    std::string logname = opt.dbdir + "/dnskv.log";
    log_fd_ = open(logname.c_str(), O_RDWR | O_CREAT, 0755);
    if(log_fd_ < 3) {
        perror("open log file error");
        exit(-1);
    }
  
    uint64_t fs = kvalloc::file_size(logname.c_str());
    log_offset_ = std::max(fs, (uint64_t)4096);
    reclaim_offset_ = log_offset_;
    db_fd_ = allocator_->file_.GetFD();
    hot_num_.store(0);
    max_hot_num_ = opt.dramsize / opt.avgsize;

    global_ending_ = false;
    #ifdef BG_RECLAIM
        // start a background log recliaim thread
        reclaim_thread_ = std::move(std::thread(&DNSKV::LogReclaim, this));
    #endif
    
    // memory reclaimation
    ebr = EpochBasedMemoryReclamationStrategy::getInstance();
}

DNSKV::~DNSKV() {
    global_ending_ = true;
    #ifdef BG_RECLAIM
        reclaim_thread_.join(); //等待回收线程结束
    #else
        LogReclaim();
    #endif

    auto ret = ftruncate64(log_fd_, 0);
    assert(ret == 0);

    delete hot_index_;
    delete global_index_;
    delete allocator_;
    close(log_fd_);
    ebr->clearAll();
    printf("miss:%lu\n", miss);
}

void DNSKV::Put(const _key_t k, const string & v, bool sync) {
    Serialize(PUT, k, v, local_buffer);
    Request * req = (Request *) local_buffer;
    FileAddr addr = Append(req, sync);
    global_index_->insert(k, addr.raw);
}

bool DNSKV::Update(const _key_t k, const string & v, bool sync) {
    Serialize(UPDATE, k, v, local_buffer);
    Request * req = (Request *) local_buffer;
    FileAddr addr = Append(req, sync); // append to WAL
    int vsize = v.size();

    // 如果数据在内存中，那么我们保证更新之后还在内存中
    #ifdef ENABLE_HOTINDEX
    update_retry:
    Request * kvitem = (Request *) hot_index_->lookup(k);
    if(kvitem != nullptr) {
        EpochGuard guard;
        kvitem->rwlock.WLock();
        if(kvitem->val_size >= vsize && kvitem->stale == 0) { // 原位更新
            memcpy(kvitem->keyvalue + 8, v.c_str(), vsize);
            kvitem->hot = 1;
            kvitem->val_size = vsize;
            kvitem->rwlock.UnWLock();
        } else { // 异位更新
            if(kvitem->stale == 1) { // 数据已经被异位更新
                kvitem->rwlock.UnWLock();
                goto update_retry;
            }
            // 新建一个新的数据项，写入到内存中
            Request * new_item = (Request *)new char[req->Length()];
            memcpy(new_item, local_buffer, req->Length());
            new_item->hot = 0;   // 初始时标记hot位为0，可以被替换，如果再次被访问，才会被标记为1（不可替换）
            new_item->stale = 0;
            hot_index_->update(k, (uint64_t)new_item); // 更新内存索引

            kvitem->stale = 1; // 由于要异位更新，需要标记数据为过期
            kvitem->hot = 0; // 将其标记为冷，之后EvictBatch的时候被替换出去
            kvitem->rwlock.UnWLock();
            
            // 插入链表
            hot_items_.push_front(new_item);
            hot_num_.fetch_add(1);

            // 当前内存空间不足，需要某个线程负责换出冷的数据
            if(hot_num_.load(std::memory_order_relaxed) >= max_hot_num_ && evict_mu_.TryLock()) {
                // 赢家来负责批量释放空间
                EvictBatch();
                evict_mu_.UnLock();
            }
        }
    }
    #endif // ENABLE_HOTINDEX

    // update global index
    FileAddr old_addr; 
    old_addr.raw = global_index_->update(k, addr.raw);

    if(old_addr.raw == 0) // not found in database 
        return false;

    if(old_addr.slabing == 1) { // free slot in kv allocator
        allocator_->Free(old_addr);
    }
    return true;
}

bool DNSKV::Delete(const _key_t k, bool sync) {
    Serialize(PUT, k, "", local_buffer);
    Request * req = (Request *) local_buffer;
    Append(req, sync);
    hot_index_->remove(k);

    FileAddr old_addr; old_addr.raw = global_index_->remove(k);
    if(old_addr.raw == 0) // not found in database 
        return false;
    
    if(old_addr.slabing == 1) { // free slot in kv allocator
        allocator_->Free(old_addr);
    }
    return true;
}

bool DNSKV::Get(const _key_t k, string & v) {
    EpochGuard guard; // 开始访问kvitem
    bool valid;

    // 优先从内存中读取
    #ifdef ENABLE_HOTINDEX
    get_retry:
    Request * kvitem = (Request *) hot_index_->lookup(k);
    if(kvitem == nullptr) {
        miss += 1;
    #endif // ENABLE_HOTINDEX
        FileAddr addr;
        addr.raw = global_index_->lookup(k); // 通过全局索引找到数据
        if(addr.warm == 0) { // 如果该数据为冷
            FileAddr new_addr = addr;
            new_addr.warm = 1;
            global_index_->update(k, new_addr.raw); // 在全局索引中标记该数据为温
        }

        if(addr.raw == 0) 
            return false; // not found in global index

        if(addr.slabing == 1) { // this value is in slab 
            allocator_->Read(local_buffer, addr);
        } else { // this value is in log
            assert(addr.FileOff() < log_offset_);
            auto discard = pread(log_fd_, local_buffer, kvalloc::MAX_SIZECLASS, addr.FileOff());
        }
        valid = Deserialize(local_buffer, v);

        #ifdef ENABLE_HOTINDEX
        Request * req = (Request *) local_buffer;
        if(valid && addr.warm == 1)   { // 该数据已经为温热，则将其Promote到热缓存中
            Request * new_item = (Request *)new char[req->Length()];
            memcpy(new_item, local_buffer, req->Length());
            new_item->hot = 0;   // 初始时标记hot位为0，可以被替换，如果再次被访问，才会被标记为1（不可替换）
            new_item->stale = 0;
            hot_index_->insert(k, (uint64_t)new_item); // 插入到内存索引
            
            // 插入链表
            hot_items_.push_front(new_item);
            hot_num_.fetch_add(1);

            // 当前内存空间不足，需要某个线程负责换出冷的数据
            if(hot_num_.load(std::memory_order_relaxed) >= max_hot_num_ && evict_mu_.TryLock()) {
                // 赢家来负责批量释放空间
                EvictBatch();
                evict_mu_.UnLock();
            }
        }
    } else {
        kvitem->rwlock.RLock();
        if(kvitem->stale == 1) { // 已经被其他线程异位更新
            kvitem->rwlock.UnRLock();
            goto get_retry;
        }

        if(kvitem->hot == 0) { // 数据被再次访问，标记为热
            kvitem->hot = 1; 
        }
        valid = Deserialize((char *)kvitem, v);
        kvitem->rwlock.UnRLock();
    }
    #endif // ENABLE_HOTINDEX
    return valid;
}

int DNSKV::Scan(const _key_t k, int len, vector<string> & out) {
    EpochGuard guard; // 开始访问kvitem
    // 从全局索引中执行范围查询
    vector<tlbtree::Record> full_res;
    global_index_->scan(k, len, full_res);

    #ifdef ENABLE_HOTINDEX
    // 从热缓存中执行范围查询
    int hot_pos = 0;
    vector<morphtree::Record> hot_res;
    int hot_size = hot_index_->scan(k, full_res.back().key, hot_res); 
    #endif // ENABLE_HOTINDEX

    IOuring ring(db_fd_, full_res.size());
    io_uring_cqe* cqe;
    // 回执参数体:当每一个异步io结束的时候，需要修改全局索引，参数存放在这里
    struct CBData {
        bool promote;
        _key_t k;
        void * buf;
    };
    vector<CBData> cbds(full_res.size());
    int used = 0;

    FileAddr addr;
    for(int i = 0; i < full_res.size(); i++) {
        out.push_back(string());
        #ifdef ENABLE_HOTINDEX
        if(hot_pos < hot_size && full_res[i].key == hot_res[hot_pos].key) {
            Request * kvitem = (Request *)hot_res[hot_pos].val;
            kvitem->rwlock.RLock();
            if(kvitem->hot == 0) { // 数据被再次访问，标记为热
                kvitem->hot = 1; 
            }
            Deserialize((char *)kvitem, out.back());
            kvitem->rwlock.UnRLock();
            hot_pos += 1;
        } else {
        #endif // ENABLE_HOTINDEX
            addr.raw = (uint64_t)full_res[i].val;
            if(addr.warm == 0) { // 如果该数据为冷
                FileAddr new_addr = addr;
                new_addr.warm = 1;
                global_index_->update(full_res[i].key, new_addr.raw); // 在全局索引中标记该数据为温
            }
            // 执行IO请求
            if(addr.slabing == 1) { // this value is in slab 
                int objsize = allocator_->GetSize(addr);
                cbds[used].promote = (addr.warm == 1);
                cbds[used].buf = new char[objsize];
                cbds[used].k = full_res[i].key;
                ring.Read(cbds[used].buf, objsize, addr.FileOff(), &cbds[used]);
                used += 1;
            } else { // this value is in log
                assert(addr.FileOff() < log_offset_); 
                auto discard = pread(log_fd_, local_buffer, kvalloc::MAX_SIZECLASS, addr.FileOff());
                bool valid = Deserialize(local_buffer, out.back());
            }

            // 防止热缓存scan错误，跳过冷热scan不匹配的结果
            while(hot_pos < hot_size && full_res[i].key > hot_res[hot_pos].key) {
                hot_pos += 1;
            }
        }

        if(used > 0) ring.Submit();

        // 等待所有io完成
        while(!ring.Empty()) {
            // 处理一个请求
            auto ret = ring.Wait(&cqe);
            assert(ret >= 0);
            // call back
            CBData * data = (CBData *)io_uring_cqe_get_data(cqe);
            ring.Seen(cqe);

            #ifdef ENABLE_HOTINDEX
            // 将scan的数据升级到热缓存中
            Request * req = (Request *) data->buf;
            out.back() = std::string((char *)data->buf, req->Length());
            if(data->promote) { // 该数据已经为温热，则将其Promote到热缓存中
                req->hot = 0;   // 初始时标记hot位为0，可以被替换，如果再次被访问，才会被标记为1（不可替换）
                req->stale = 0;
                hot_index_->insert(data->k, (uint64_t)data->buf); // 插入到内存索引
                
                // 插入链表
                hot_items_.push_front(req);
                hot_num_.fetch_add(1);
            } else {
                delete [] ((char *)data->buf);
            }
            #endif // ENABLE_HOTINDEX
        }
    }

    #ifdef ENABLE_HOTINDEX
    // 当前内存空间不足，需要某个线程负责换出冷的数据
    if(hot_num_.load(std::memory_order_relaxed) >= max_hot_num_ && evict_mu_.TryLock()) {
        // 赢家来负责批量释放空间
        EvictBatch(hot_num_.load() - max_hot_num_);
        evict_mu_.UnLock();
    }
    #endif // ENABLE_HOTINDEX

    return out.size();
}

FileAddr DNSKV::Append(Request * req, bool sync) {
    uint64_t addr = __sync_fetch_and_add(&log_offset_, req->Length());
    auto discard = pwrite(log_fd_, req, req->Length(), addr);
    if(sync) fsync(log_fd_);
    FileAddr res;
    res.slabing = 0;
    res.warm = 1; // 写入新数据被认为是温热的
    res.slabid = addr / kvalloc::SLAB_SIZE;
    res.slaboff = addr % kvalloc::SLAB_SIZE;
    return res;
}

void DNSKV::EvictBatch(int num) {
    static const int EVICT_BATCH = 20;
    static const int MAX_PROBE = 1000;

    if(num == 0) 
        num = EVICT_BATCH;
    
    // 一次性替换多个冷数据
    for(int i = 0; i < num; i++) {
        uint32_t probe_len = 0;
        Request * curitem = hot_items_.get_next(); // 定位到clock环的第一个
        while(curitem->hot == 1 && probe_len < MAX_PROBE) {
            probe_len += 1;
            curitem->hot = 0; // 数据原来为1，现在降温
            curitem = hot_items_.go_next(); // 跳转到clock环的下一个
        } 

        hot_items_.erase_next(); // 下一个温度为0， 删除它
        hot_num_.fetch_sub(1);
        // 第一步: 从局部索引中删除，并且删除内存空间
        _key_t k = *((uint64_t *)(curitem->keyvalue));
        hot_index_->remove(k);
        
        // 第二步：在全局索引中标记该数据为冷
        FileAddr addr;
        addr.raw = global_index_->lookup(k); 
        addr.warm = 0; 
        global_index_->update(k, addr.raw);

        // 第三步：释放内存空间
        ebr->scheduleForDeletion((char *)curitem);
    }
}

static constexpr uint32_t RECLAIM_BATCH = 8 * 1024 * 1024; // 8 MiB
char RECLAIM_BUFFER[RECLAIM_BATCH];
void DNSKV::LogReclaim() {
    static constexpr int QD = 128;
    IOuring ring(db_fd_, QD);
    io_uring_cqe* cqe;
    // 回执参数体:当每一个异步io结束的时候，需要修改全局索引，参数存放在这里
    struct CBData {
        _key_t k;
        FileAddr addr;
    };
    // 预先分配所有需要使用的CBData
    CBData cbdatas[QD];
    std::deque<CBData *> free_list_;
    for(int i = 0; i < QD; i++) {
        free_list_.push_back(&cbdatas[i]);
    }

    // 死循环处理日志数据，除非日志全部处理完且全局标记结束
    while(!(global_ending_ && log_offset_ == reclaim_offset_)) {
        // 需要保证log_offset_比reclaim_offset_足够前，防止数据还没写入log文件
        if(log_offset_ - reclaim_offset_ >= RECLAIM_BATCH * 2 || global_ending_) { // 现在log文件中有足够多的日志数据，或者现在已经标记全局结束
            // 从log文件中读取一大块日志区域
            size_t read_bytes = pread(log_fd_, RECLAIM_BUFFER, RECLAIM_BATCH, reclaim_offset_);
            uint32_t inbuffer_off = 0;
            bool reach_end = false;

            // 依次将所有日志区域写入ring
            while (inbuffer_off < read_bytes && reach_end == false) { // 判断是否到达一个batch的尾部
                // 先尽量填满一个uring
                while(!ring.Full() && inbuffer_off < read_bytes) {
                    Request * req = (Request *) &RECLAIM_BUFFER[inbuffer_off];
                    ReqType op;
                    _key_t k;
                    uint32_t req_len;

                    if(read_bytes - inbuffer_off <= sizeof(Request) || read_bytes - inbuffer_off < req->Length()) {
                        // 已经到达文件的末尾
                        reach_end = true;
                        break;
                    } else {
                        op = (ReqType) req->op;
                        k = *((_key_t *)(req->keyvalue));
                        req_len = req->Length();
                        FileAddr latest_addr; 
                        latest_addr.raw = global_index_->lookup(k);
                        assert(latest_addr.raw > 0);
                        if(op != DELETE && latest_addr.FileOff() == reclaim_offset_ + inbuffer_off) { // 该处能够跳过陈旧的日志项
                            FileAddr new_addr = allocator_->Alloc(req_len);
                            new_addr.slabing = 1;
                            new_addr.warm = 0;
                            // 留下一个回执参数体，方便后续更新全局索引
                            CBData * data = free_list_.front(); 
                            free_list_.pop_front();
                            data->k = k;
                            data->addr = new_addr;
                            ring.Write(RECLAIM_BUFFER + inbuffer_off, req_len, new_addr.FileOff(), data);
                        }
                        inbuffer_off += req_len; // 处理完毕一个日志项
                    }
                } 
                // 将ring提交，如果为空，则直接重启循环
                if(ring.Empty()) {
                    continue;
                } else {
                    ring.Submit();
                }
                
                // 等待并处理一个请求
                auto ret = ring.Wait(&cqe);
                if(ret < 0) {
                    int a = 0;
                }
                // call back
                CBData * data = (CBData *)io_uring_cqe_get_data(cqe);
                auto old = global_index_->update(data->k, data->addr.raw);
                free_list_.push_back(data);
                ring.Seen(cqe);
            }

            // 等待所有异步IO结束，回收之前的全部空间
            while(!ring.Empty()) {
                // 处理一个请求
                auto ret = ring.Wait(&cqe);
                if(ret < 0) {
                    int a = 0;
                }
                // call back
                CBData * data = (CBData *)io_uring_cqe_get_data(cqe);
                auto old = global_index_->update(data->k, data->addr.raw);
                assert(old > 0);
                free_list_.push_back(data);
                ring.Seen(cqe);
            }
            
            // 确定完成数据的回收，回收日志文件空间
            auto ret = fallocate64(log_fd_, FALLOC_FL_PUNCH_HOLE | FALLOC_FL_KEEP_SIZE, 
                                        reclaim_offset_, inbuffer_off); // 文件打洞，删除对应的块内容
            assert(ret == 0);
            reclaim_offset_ += inbuffer_off;
        } else { // 等待5ms足够多的数据进入日志文件
            std::this_thread::sleep_for(std::chrono::milliseconds(5)); 
        }
    }
    assert(reclaim_offset_ == log_offset_);
}

} // namespace dnskv